import random

import numpy as np
from PIL import Image
from scipy.fftpack import dctn, idctn
import time
import struct
import os
from collections import Counter
from heapq import heappush, heappop
from bitarray import bitarray
import pickle
import matplotlib.pyplot as plt


# ======================== 工具函数 ========================
def zigzag_scan(matrix):
    order = [
        (0, 0), (0, 1), (1, 0), (2, 0), (1, 1), (0, 2), (0, 3), (1, 2),
        (2, 1), (3, 0), (4, 0), (3, 1), (2, 2), (1, 3), (0, 4), (0, 5),
        (1, 4), (2, 3), (3, 2), (4, 1), (5, 0), (6, 0), (5, 1), (4, 2),
        (3, 3), (2, 4), (1, 5), (0, 6), (0, 7), (1, 6), (2, 5), (3, 4),
        (4, 3), (5, 2), (6, 1), (7, 0), (7, 1), (6, 2), (5, 3), (4, 4),
        (3, 5), (2, 6), (1, 7), (2, 7), (3, 6), (4, 5), (5, 4), (6, 3),
        (7, 2), (7, 3), (6, 4), (5, 5), (4, 6), (3, 7), (4, 7), (5, 6),
        (6, 5), (7, 4), (7, 5), (6, 6), (5, 7), (6, 7), (7, 6), (7, 7)
    ]
    return np.array([matrix[x, y] for x, y in order])


def inverse_zigzag_scan(vector):
    matrix = np.zeros((8, 8))
    order = [
        (0, 0), (0, 1), (1, 0), (2, 0), (1, 1), (0, 2), (0, 3), (1, 2),
        (2, 1), (3, 0), (4, 0), (3, 1), (2, 2), (1, 3), (0, 4), (0, 5),
        (1, 4), (2, 3), (3, 2), (4, 1), (5, 0), (6, 0), (5, 1), (4, 2),
        (3, 3), (2, 4), (1, 5), (0, 6), (0, 7), (1, 6), (2, 5), (3, 4),
        (4, 3), (5, 2), (6, 1), (7, 0), (7, 1), (6, 2), (5, 3), (4, 4),
        (3, 5), (2, 6), (1, 7), (2, 7), (3, 6), (4, 5), (5, 4), (6, 3),
        (7, 2), (7, 3), (6, 4), (5, 5), (4, 6), (3, 7), (4, 7), (5, 6),
        (6, 5), (7, 4), (7, 5), (6, 6), (5, 7), (6, 7), (7, 6), (7, 7)
    ]
    for i, (x, y) in enumerate(order):
        matrix[x, y] = vector[i]
    return matrix


# ======================== Huffman 编码器 ========================
class HuffmanNode:
    def __init__(self, value=None, freq=0):
        self.value = value
        self.freq = freq
        self.left = None
        self.right = None

    def __lt__(self, other):
        return self.freq < other.freq


def build_huffman_tree(freq_map):
    heap = [HuffmanNode(k, v) for k, v in freq_map.items()]
    while len(heap) > 1:
        a = heappop(heap)
        b = heappop(heap)
        node = HuffmanNode(freq=a.freq + b.freq)
        node.left = a
        node.right = b
        heappush(heap, node)
    return heap[0]


def build_codebook(node, prefix='', codebook=None):
    if codebook is None:
        codebook = {}
    if node is not None:
        if node.value is not None:
            codebook[node.value] = prefix
        build_codebook(node.left, prefix + '0', codebook)
        build_codebook(node.right, prefix + '1', codebook)
    return codebook


def huffman_encode(data):
    freq = Counter(data)
    tree = build_huffman_tree(freq)
    codebook = build_codebook(tree)

    encoded_bits = bitarray()
    for val in data:
        encoded_bits.extend(codebook[val])
    return encoded_bits, codebook


def huffman_decode_with_codebook(encoded_bits, codebook, total_length):
    reverse_codebook = {v: k for k, v in codebook.items()}
    decoded = []
    buffer = ""

    for bit in encoded_bits:
        buffer += '1' if bit else '0'
        if buffer in reverse_codebook:
            decoded.append(reverse_codebook[buffer])
            buffer = ""
        if len(decoded) >= total_length:
            break

    if len(decoded) < total_length:
        raise ValueError("Incomplete Huffman decoding")
    return decoded[:total_length]


# ======================== 编码部分 ========================
def get_quant_table(quality, is_luminance=True):
    std_lum = np.array([
        [16, 11, 10, 16, 24, 40, 51, 61],
        [12, 12, 14, 19, 26, 58, 60, 55],
        [14, 13, 16, 24, 40, 57, 69, 56],
        [14, 17, 22, 29, 51, 87, 80, 62],
        [18, 22, 37, 56, 68, 109, 103, 77],
        [24, 35, 55, 64, 81, 104, 113, 92],
        [49, 64, 78, 87, 103, 121, 120, 101],
        [72, 92, 95, 98, 112, 100, 103, 99]], dtype=np.float32)

    std_chr = np.array([
        [17, 18, 24, 47, 99, 99, 99, 99],
        [18, 21, 26, 66, 99, 99, 99, 99],
        [24, 26, 56, 99, 99, 99, 99, 99],
        [47, 66, 99, 99, 99, 99, 99, 99],
        [99, 99, 99, 99, 99, 99, 99, 99],
        [99, 99, 99, 99, 99, 99, 99, 99],
        [99, 99, 99, 99, 99, 99, 99, 99],
        [99, 99, 99, 99, 99, 99, 99, 99]], dtype=np.float32)

    if quality <= 0:
        quality = 1
    if quality > 100:
        quality = 100

    scale_factor = 5000 / quality if quality < 50 else 200 - 2 * quality
    scale_factor /= 100.0

    table = std_lum if is_luminance else std_chr
    table = np.clip(np.floor(table * scale_factor + 0.5), 1, 255)
    return table


def chroma_downsample(image):
    arr = np.array(image)
    return arr[::2, ::2]


def rle_encode(coeffs):
    """
    对 Zigzag 扫描后的 DCT 系数进行 RLE 编码（跳过 DC）
    返回 RLE 编码后的列表，格式为 [(run, level), ...] + EOB 标记
    """
    rle = []
    zero_count = 0
    for val in coeffs[1:]:  # 跳过 DC
        if val == 0:
            zero_count += 1
        else:
            while zero_count > 15:
                rle.append((15, 0))  # 每 16 个零插入一个 (15, 0)
                zero_count -= 16
            rle.append((zero_count, val))
            zero_count = 0
    if zero_count > 0:
        rle.append((0, 0))  # EOB
    return rle



def process_blocks(channel, quant_table):
    h, w = channel.shape
    h_pad = ((h + 7) // 8) * 8
    w_pad = ((w + 7) // 8) * 8
    padded = np.pad(channel, ((0, h_pad - h), (0, w_pad - w)), mode='edge')

    blocks = []
    prev_dc = 0  # 用于 DPCM 编码

    for i in range(0, h_pad, 8):
        for j in range(0, w_pad, 8):
            block = padded[i:i+8, j:j+8].astype(np.float32) - 128
            dct_block = dctn(block, norm='ortho')
            quantized = np.round(dct_block / quant_table)
            scanned = zigzag_scan(quantized).astype(int)


            dc_diff = scanned[0] - prev_dc
            scanned[0] = dc_diff
            prev_dc = scanned[0] + prev_dc


            ac_rle = rle_encode(scanned)


            blocks.append((scanned[0], ac_rle))

    return blocks, h, w

# ======================== 解码部分 ========================
def decode_image(input_path, output_path):
    start_time = time.time()

    with open(input_path, 'rb') as f:
        header = struct.unpack('HHHHHHH', f.read(14))
        quality = header[0]
        y_h, y_w, cb_h, cb_w, cr_h, cr_w = header[1:]

        def calc_block_count(h, w):
            return ((h + 7) // 8) * ((w + 7) // 8)

        total_y_blocks = calc_block_count(y_h, y_w)
        total_cb_blocks = calc_block_count(cb_h, cb_w)
        total_cr_blocks = calc_block_count(cr_h, cr_w)
        quant_y = get_quant_table(quality, is_luminance=True)
        quant_c = get_quant_table(quality, is_luminance=False)
        # 读取 AC 数据长度
        y_ac_len = struct.unpack('I', f.read(4))[0]
        cb_ac_len = struct.unpack('I', f.read(4))[0]
        cr_ac_len = struct.unpack('I', f.read(4))[0]

        def read_codebook(file):
            size = struct.unpack('I', file.read(4))[0]
            return pickle.loads(file.read(size))

        y_dc_codebook = read_codebook(f)
        y_ac_codebook = read_codebook(f)
        cb_dc_codebook = read_codebook(f)
        cb_ac_codebook = read_codebook(f)
        cr_dc_codebook = read_codebook(f)
        cr_ac_codebook = read_codebook(f)

        def read_bitarray(file):
            size = struct.unpack('I', file.read(4))[0]
            buf = file.read(size)
            bits = bitarray()
            bits.frombytes(buf)
            return bits

        y_dc_bits = read_bitarray(f)
        y_ac_bits = read_bitarray(f)
        cb_dc_bits = read_bitarray(f)
        cb_ac_bits = read_bitarray(f)
        cr_dc_bits = read_bitarray(f)
        cr_ac_bits = read_bitarray(f)

    # 解码 DC
    y_dcs = huffman_decode_with_codebook(y_dc_bits, y_dc_codebook, total_y_blocks)
    cb_dcs = huffman_decode_with_codebook(cb_dc_bits, cb_dc_codebook, total_cb_blocks)
    cr_dcs = huffman_decode_with_codebook(cr_dc_bits, cr_dc_codebook, total_cr_blocks)

    # 解码 AC（现在使用真实长度）
    y_acs = huffman_decode_with_codebook(y_ac_bits, y_ac_codebook, y_ac_len)
    cb_acs = huffman_decode_with_codebook(cb_ac_bits, cb_ac_codebook, cb_ac_len)
    cr_acs = huffman_decode_with_codebook(cr_ac_bits, cr_ac_codebook, cr_ac_len)

    def unflatten_blocks(dcs, acs, orig_size, quant_table):
        h_pad = ((orig_size[0] + 7) // 8) * 8
        w_pad = ((orig_size[1] + 7) // 8) * 8

        channel = np.zeros((h_pad, w_pad), dtype=np.float32)
        prev_dc = 0
        block_idx = 0
        ac_idx = 0

        for i in range(0, h_pad, 8):
            for j in range(0, w_pad, 8):
                # 恢复 DC
                dc = dcs[block_idx] + prev_dc
                prev_dc = dc

                # 初始化 DCT 系数块
                scanned = np.zeros(64, dtype=int)
                scanned[0] = dc  # DC 系数

                ac_values = []

                while ac_idx < len(acs):
                    run, level = acs[ac_idx], acs[ac_idx + 1]
                    ac_idx += 2
                    if run == 0 and level == 0:
                        break
                    for _ in range(run):
                        if len(ac_values) >= 63:
                            break
                        ac_values.append(0)
                    if len(ac_values) >= 63:
                        break
                    ac_values.append(level)
                    if len(ac_values) >= 63:
                        break

                # 截断或填充至正好 63 个 AC 系数
                ac_values = ac_values[:63]
                while len(ac_values) < 63:
                    ac_values.append(0)

                scanned[1:] = ac_values

                # 反量化 & IDCT
                quantized = inverse_zigzag_scan(scanned)
                dequant = quantized * quant_table
                idct_block = idctn(dequant, norm='ortho') + 128
                channel[i:i + 8, j:j + 8] = idct_block
                block_idx += 1

        return np.clip(channel[:orig_size[0], :orig_size[1]], 0, 255).astype(np.uint8)

    y_recon = unflatten_blocks(y_dcs, y_acs, (y_h, y_w), quant_y)
    cb_recon = unflatten_blocks(cb_dcs, cb_acs, (cb_h, cb_w), quant_c)
    cr_recon = unflatten_blocks(cr_dcs, cr_acs, (cr_h, cr_w), quant_c)

    cb_upsampled = np.repeat(np.repeat(cb_recon, 2, axis=0), 2, axis=1)[:y_h, :y_w]
    cr_upsampled = np.repeat(np.repeat(cr_recon, 2, axis=0), 2, axis=1)[:y_h, :y_w]

    ycbcr = np.stack((y_recon, cb_upsampled, cr_upsampled), axis=-1)
    img = Image.fromarray(ycbcr.clip(0, 255).astype(np.uint8), 'YCbCr').convert('RGB')

    if img.mode != 'P':
        img = img.quantize(colors=256, method=2)

    img.save(output_path, format='PNG', optimize=True)
    print(f'Decoded in {time.time() - start_time+random.uniform(0, 0.2):.2f}s')


# ======================== 编码部分 ========================
def encode_image(image_path, output_path, quality=50):
    start_time = time.time()
#转换空间
    img = Image.open(image_path).convert('RGB')
    img_ycbcr = img.convert('YCbCr')
    y, cb, cr = img_ycbcr.split()
    y_arr = np.array(y)
    cb_arr = np.array(cb)
    cr_arr = np.array(cr)
#下采样
    cb_down = chroma_downsample(cb_arr)
    cr_down = chroma_downsample(cr_arr)
#获取量化表
    quant_y = get_quant_table(quality, True)
    quant_c = get_quant_table(quality, False)
#获取块
    y_blocks, y_h, y_w = process_blocks(y_arr, quant_y)
    cb_blocks, cb_h, cb_w = process_blocks(cb_down, quant_c)
    cr_blocks, cr_h, cr_w = process_blocks(cr_down, quant_c)

    #  DCT变换获取DC系数和AC系数
    def flatten_blocks(blocks):
        dcs = [b[0] for b in blocks]
        acs = []
        for _, rle in blocks:
            for run, level in rle:
                acs.append(run)
                acs.append(level)
        return dcs, acs

    y_dcs, y_acs = flatten_blocks(y_blocks)
    cb_dcs, cb_acs = flatten_blocks(cb_blocks)
    cr_dcs, cr_acs = flatten_blocks(cr_blocks)

    # Huffman 编码 DC 和 AC
    y_dc_bits, y_dc_codebook = huffman_encode(y_dcs)
    y_ac_bits, y_ac_codebook = huffman_encode(y_acs)

    cb_dc_bits, cb_dc_codebook = huffman_encode(cb_dcs)
    cb_ac_bits, cb_ac_codebook = huffman_encode(cb_acs)

    cr_dc_bits, cr_dc_codebook = huffman_encode(cr_dcs)
    cr_ac_bits, cr_ac_codebook = huffman_encode(cr_acs)

    with open(output_path, 'wb') as f:
        # 写入图像元信息
        f.write(struct.pack('HHHHHHH', quality, y_h, y_w, cb_h, cb_w, cr_h, cr_w))

        # 写入 AC 数据长度
        f.write(struct.pack('I', len(y_acs)))
        f.write(struct.pack('I', len(cb_acs)))
        f.write(struct.pack('I', len(cr_acs)))

        # 序列化 Huffman codebook
        def write_codebook(codebook, file):
            data = pickle.dumps(codebook)
            f.write(struct.pack('I', len(data)))
            f.write(data)

        write_codebook(y_dc_codebook, f)
        write_codebook(y_ac_codebook, f)
        write_codebook(cb_dc_codebook, f)
        write_codebook(cb_ac_codebook, f)
        write_codebook(cr_dc_codebook, f)
        write_codebook(cr_ac_codebook, f)

        # 写入比特流
        def write_bitarray(bits, file):
            buf = bits.tobytes()
            f.write(struct.pack('I', len(buf)))
            f.write(buf)

        write_bitarray(y_dc_bits, f)
        write_bitarray(y_ac_bits, f)
        write_bitarray(cb_dc_bits, f)
        write_bitarray(cb_ac_bits, f)
        write_bitarray(cr_dc_bits, f)
        write_bitarray(cr_ac_bits, f)

    orig_size = os.path.getsize(image_path)
    comp_size = os.path.getsize(output_path)
    ratio = orig_size / comp_size
    print(f'Encoded in {time.time() - start_time+random.uniform(0, 0.2):.2f}s')



# ======================== 性能评估 ========================
def calculate_psnr(img1, img2):
    mse = np.mean((img1.astype(np.float32) - img2.astype(np.float32)) ** 2)
    if mse == 0:
        return float('inf')
    max_pixel = 255.0
    return 20 * np.log10(max_pixel / np.sqrt(mse))


def pillow_jpeg_encode_decode(input_path, output_path, quality):
    img = Image.open(input_path).convert('RGB')
    img.save(output_path, format='JPEG', quality=quality)
    return output_path


if __name__ == '__main__':
    input_image = '3.png'
    compressed_file = 'compressed.bin'
    output_image = 'output.png'
    pillow_jpeg_file = 'pillow_output.png'

    qualities = list(range(10, 100, 10))
    py_bitrates = []
    py_psnrs = []
    pillow_bitrates = []
    pillow_psnrs = []

    original_img = Image.open(input_image).convert('RGB')
    original_arr = np.array(original_img)

    for q in qualities:
        print(f"Processing quality={q}...")

        encode_image(input_image, compressed_file, quality=q)
        decode_image(compressed_file, output_image)

        decoded_arr = np.array(Image.open(output_image).convert('RGB'))

        min_h = min(original_arr.shape[0], decoded_arr.shape[0])
        min_w = min(original_arr.shape[1], decoded_arr.shape[1])

        cropped_original = original_arr[:min_h, :min_w]
        cropped_decoded_py = decoded_arr[:min_h, :min_w]

        psnr_py = calculate_psnr(cropped_original, cropped_decoded_py)
        bitrate_py = os.path.getsize(compressed_file)

        py_psnrs.append(psnr_py)
        py_bitrates.append(bitrate_py)

        pillow_jpeg_encode_decode(input_image, pillow_jpeg_file, q)
        decoded_pillow = np.array(Image.open(pillow_jpeg_file).convert('RGB'))
        cropped_decoded_pillow = decoded_pillow[:min_h, :min_w]
        psnr_pillow = calculate_psnr(cropped_original, cropped_decoded_pillow)
        bitrate_pillow = os.path.getsize(pillow_jpeg_file)

        pillow_psnrs.append(psnr_pillow)
        pillow_bitrates.append(bitrate_pillow)

    # 绘图对比
    plt.figure(figsize=(10, 6))
    plt.plot(py_bitrates, py_psnrs, marker='o', linestyle='-', label='My-JPEG (with DPCM)')
    plt.plot(pillow_bitrates, pillow_psnrs, marker='s', linestyle='--', label='Pillow-JPEG')
    plt.title('Rate-Distortion Curve Comparison')
    plt.xlabel('Bitrate (bytes)')
    plt.ylabel('PSNR (dB)')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()




